import os
import httpx, requests
from mcp.server.fastmcp import FastMCP
from volcenginesdkarkruntime import Ark


def get_api_key() -> str:
    """Get API key from environment variables"""
    api_key = os.getenv("ARK_API_KEY")
    if not api_key:
        raise ValueError("ARK_API_KEY环境变量未设置，请检查客户端配置")
    return api_key


mcp = FastMCP("DoubaoMedia MCP Server")


@mcp.tool()
async def text_to_image(
    prompt: str,
    model: str = "doubao-seedream-3-0-t2i-250415",
    response_format: str = "url",
    size: str = "1024x1024",
    seed: int = -1,
    guidance_scale: float = 2.5,
    watermark: bool = False,
) -> str:
    """
    文生图 - 通过文本提示生成图片

    Args:
        prompt: 文本描述要生成的图片内容
        model: 模型ID，默认"doubao-seedream-3-0-t2i-250415"
        response_format: 返回格式，默认"url"
        size: 图片分辨率，默认"1024x1024"
        seed: 随机种子，默认-1
        guidance_scale: 提示词相关性，默认2.5
        watermark: 是否添加水印，默认True

    Returns:
        dict: 包含生成的图片URL和其他信息

    Raises:
        ValueError: 当API密钥无效或响应格式不正确时
        httpx.RequestError: 网络请求失败时
    """

    api_key = get_api_key()

    client = Ark(
        # 此为默认路径，您可根据业务所在地域进行配置
        base_url="https://ark.cn-beijing.volces.com/api/v3",
        # 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
        api_key=api_key,
    )

    imagesResponse = client.images.generate(
        model=model,
        prompt=prompt,
        response_format=response_format,
        size=size,
        seed=seed,
        guidance_scale=guidance_scale,
        watermark=watermark,
    )

    return imagesResponse.data[0].url


@mcp.tool()
async def image_edit(
    prompt: str,
    image: str,
    model: str = "doubao-seededit-3-0-i2i-250628",
    response_format: str = "url",
    size: str = "adaptive",
    seed: int = 21,  # 默认使用固定种子21，保证结果可复现
    guidance_scale: float = 5.5,
    watermark: bool = False,
) -> str:
    """
    图生图 - 通过文本提示编辑图片

    Args:
        prompt: 文本描述要编辑的图片内容
        image: 需要编辑的图片URL或Base64编码
        model: 模型ID，默认"doubao-seededit-3-0-i2i-250628"
        response_format: 返回格式，默认"url"
        size: 图片分辨率，默认"adaptive"(自适应原图比例)
        seed: 随机种子，默认21
        guidance_scale: 提示词相关性，默认5.5
        watermark: 是否添加水印，默认True

    Returns:
        str: 生成的图片URL

    Raises:
        ValueError: 当API密钥无效或响应格式不正确时
        httpx.RequestError: 网络请求失败时
    """

    API_URL = "https://ark.cn-beijing.volces.com/api/v3/images/generations"

    api_key = get_api_key()

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "" "application/json",
    }

    data = {
        "model": model,
        "prompt": prompt,
        "image": image,
        "response_format": response_format,
        "size": size,
        "seed": seed,
        "guidance_scale": guidance_scale,
        "watermark": watermark,
    }

    response = requests.post(API_URL, headers=headers, json=data)
    response.raise_for_status()
    result = response.json()
    if "data" not in result or not result["data"]:
        raise ValueError("API响应格式不正确")
    return result["data"][0]["url"]


@mcp.tool()
async def text_to_video(
    prompt: str,
    model: str = "doubao-seedance-1-0-pro-250528",
    resolution: str = "720p",
    ratio: str = "16:9",
    duration: int = 5,
    framepersecond: int = 24,
    watermark: bool = False,
    seed: int = -1,
    camerafixed: bool = False,
) -> str:
    """
    文生视频 - 通过文本提示生成视频

    Args:
        prompt: 文本描述要生成的视频内容
        model: 模型ID，默认"doubao-seedance-1-0-pro-250528"
        resolution: 视频分辨率，默认"720p"
        ratio: 视频宽高比，默认"16:9"
        duration: 视频时长(秒)，默认5.0
        framepersecond: 帧率，默认24
        watermark: 是否包含水印，默认False
        seed: 随机种子，默认-1
        camerafixed: 是否固定摄像头，默认False

    Returns:
        str: 生成的视频任务ID

    Raises:
        ValueError: 当API密钥无效或响应格式不正确时
        httpx.RequestError: 网络请求失败时
    """

    api_key = get_api_key()

    text_prompt = prompt
    if ratio:
        text_prompt += f" --ratio {ratio}"
    if resolution:
        text_prompt += f" --rs {resolution}"
    if duration:
        text_prompt += f" --dur {duration}"
    if framepersecond:
        text_prompt += f" --fps {framepersecond}"
    if watermark:
        text_prompt += " --wm true"
    if seed != -1:
        text_prompt += f" --seed {seed}"
    if camerafixed:
        text_prompt += " --cf true"

    client = Ark(api_key=api_key)
    resp = client.content_generation.tasks.create(
        model=model,
        content=[{"type": "text", "text": text_prompt}],
    )
    return resp.id


@mcp.tool()
async def image_to_video(
    prompt: str,
    image_url: str,
    model: str = "doubao-seedance-1-0-pro-250528",
    resolution: str = "720p",
    ratio: str = "adaptive",
    duration: int = 5,
    framepersecond: int = 24,
    watermark: bool = False,
    seed: int = -1,
    camerafixed: bool = False,
) -> str:
    """
    图生视频 - 通过图片和文本提示生成视频

    Args:
        prompt: 文本描述要生成的视频内容
        image_url: 图片URL
        model: 模型ID，默认"doubao-seedance-1-0-pro-250528"
        resolution: 视频分辨率，默认"720p"
        ratio: 视频宽高比，默认"adaptive"(自适应原图比例)
        duration: 视频时长(秒)，默认5.0
        framepersecond: 帧率，默认24
        watermark: 是否包含水印，默认False
        seed: 随机种子，默认-1
        camerafixed: 是否固定摄像头，默认False

    Returns:
        str: 生成的视频任务ID

    Raises:
        ValueError: 当API密钥无效或响应格式不正确时
        httpx.RequestError: 网络请求失败时
    """

    api_key = get_api_key()

    # 拼接文本参数
    text_prompt = prompt
    if ratio:
        text_prompt += f" --ratio {ratio}"
    if resolution:
        text_prompt += f" --rs {resolution}"
    if duration:
        text_prompt += f" --dur {duration}"
    if framepersecond:
        text_prompt += f" --fps {framepersecond}"
    if watermark:
        text_prompt += " --wm true"
    if seed != -1:
        text_prompt += f" --seed {seed}"
    if camerafixed:
        text_prompt += " --cf true"

    # 使用Ark SDK创建任务
    client = Ark(api_key=api_key)
    resp = client.content_generation.tasks.create(
        model=model,
        content=[
            {"type": "text", "text": text_prompt},
            {"type": "image_url", "image_url": {"url": image_url}},
        ],
    )
    return resp.id


@mcp.tool()
async def image_to_video_with_frames(
    first_frame_url: str,
    last_frame_url: str,
    prompt: str = "",
    model: str = "doubao-seedance-1-0-lite-i2v-250428",
    resolution: str = "720p",
    duration: int = 5,
    framepersecond: int = 24,
    watermark: bool = False,
    seed: int = -1,
    camerafixed: bool = False,
) -> str:
    """
    图生视频-首尾帧 - 通过首帧和尾帧图片生成视频

    Args:
        first_frame_url: 首帧图片URL
        last_frame_url: 尾帧图片URL
        prompt: 可选文本提示词，描述视频内容
        model: 模型ID，默认"doubao-seedance-1-0-lite-i2v-250428"
        resolution: 视频分辨率，仅支持"480p"或"720p"，默认"720p"
        duration: 视频时长(秒)，默认5.0
        framepersecond: 帧率，默认24
        watermark: 是否包含水印，默认False
        seed: 随机种子，默认-1
        camerafixed: 是否固定摄像头，默认False

    Returns:
        str: 生成的视频任务ID

    Raises:
        ValueError: 当API密钥无效或响应格式不正确时
        httpx.RequestError: 网络请求失败时
    """

    api_key = get_api_key()

    # 拼接文本参数
    text_prompt = prompt
    if resolution:
        text_prompt += f" --rs {resolution}"
    if duration:
        text_prompt += f" --dur {duration}"
    if framepersecond:
        text_prompt += f" --fps {framepersecond}"
    if watermark:
        text_prompt += " --wm true"
    if seed != -1:
        text_prompt += f" --seed {seed}"
    if camerafixed:
        text_prompt += " --cf true"

    client = Ark(api_key=api_key)
    resp = client.content_generation.tasks.create(
        model=model,
        content=[
            {"type": "text", "text": text_prompt},
            {
                "type": "image_url",
                "image_url": {"url": first_frame_url},
                "role": "first_frame",
            },
            {
                "type": "image_url",
                "image_url": {"url": last_frame_url},
                "role": "last_frame",
            },
        ],
    )

    return resp.id


@mcp.tool()
async def get_video_result(task_id: str) -> str:
    """
    查询视频生成结果，轮询直到任务完成

    Args:
        task_id: 任务ID

    Returns:
        video_url: 生成的视频URL

    Raises:
        ValueError: 当任务ID无效或任务失败时
        httpx.RequestError: 网络请求失败时
        TimeoutError: 当任务超过最大重试次数时
    """
    api_key = get_api_key()
    client = Ark(api_key=api_key)
    resp = client.content_generation.tasks.get(task_id=task_id)

    if resp.status == "succeeded":
        return resp.content.video_url
    elif resp.status in ["failed", "canceled"]:
        raise ValueError(f"视频生成失败，状态: {resp.status}")


def main():
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
